def norm(img):
img = np.array(img, dtype=np.float32)
img -= img.min()
img /= img.max()
return img
def get_annotated_img(img, annt, is_pred=False):
img = norm(np.array(np.transpose(img, (1, 2, 0)))) * 255.0 # multiple by 255 as in the dataset we divide it
# needed here because the image come from the dataset in values of [0, 1]
# and the annotations are in values of [0, 255] and plt should get images
# with values of [0, 1], therefore a normalizing of the images is needed
num_faces = len(annt['boxes'])
for i in range(num_faces):
box, cat = annt['boxes'][i], annt['labels'][i]
if is_pred:
score = annt['scores'][i]
xmin, ymin, xmax, ymax = np.array(box, dtype=np.int32)
start_point = (xmin, ymin)
end_point = (xmax, ymax)
# [None, 'without_mask','with_mask','mask_weared_incorrect']
color = (0, 0, 0)
if cat == 2:
color = (0, 255, 0) # green
elif cat == 3:
color = (0, 0, 255) # blue
elif cat == 1:
color = (255, 0, 0) # red
thickness = 1
img = cv2.rectangle(img, start_point, end_point, color, thickness)
font = cv2.FONT_HERSHEY_SIMPLEX
fontScale = 1 / 3
if not is_pred:
img = cv2.putText(img, '{}'.format(faceMasksData.classes[cat]), start_point, font, fontScale, color = (0, 0, 0), thickness=2)
img = cv2.putText(img, '{}'.format(faceMasksData.classes[cat]), start_point, font, fontScale, color, thickness)
else:
img = cv2.putText(img, '{:.2f}'.format(score), start_point, font, fontScale, color = (0, 0, 0), thickness=2)
img = cv2.putText(img, '{:.2f}'.format(score), start_point, font, fontScale, color, thickness)
return img
def show_augmented_samples(dataloader, model=None):
samples, annotations = next(iter(dataloader)) # get the first batch
amount_samples = min(12, len(samples)) # number of examples will be at most 12
for i, (img, ant) in enumerate(zip(samples, annotations)):
if i >= amount_samples:
break
plt.figure(figsize=(10, 10))
# plt.subplot(amount_samples, 1, i + 1)
img = get_annotated_img(img, ant)
plt.imshow(norm(img))
plt.axis('off')
plt.tight_layout()
plt.show()
show_augmented_samples(test_loader)
